'''
Author: didiplus
Date: 2022-12-03 16:57:35
LastEditors: didiplus
LastEditTime: 2022-12-05 21:24:13
Description: 
'''
import sys,base64,os
from Crypto import Random
from Crypto.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5
from Crypto.PublicKey import RSA


basedir = os.path.abspath(os.path.dirname(__file__))

class RsaEncryption(object):

    def __init__(self,private_file_name,public_file_name):
        self.private_file_name=private_file_name
        self.public_file_name=public_file_name
        self.length=1024
        self.random_generator = Random.new().read

    def create_rsa_key(self):
        rsa = RSA.generate(self.length,self.random_generator)
        #生成密钥对
        #生成私钥文件
        property_pem = rsa.exportKey()
        with open(os.path.join(basedir,self.private_file_name),'wb') as f:
            f.write(property_pem)
        
        #生成公钥文件
        public_pem = rsa.publickey().exportKey()
        with open(os.path.join(basedir,self.public_file_name),'wb') as f:
            f.write(public_pem)
    
    def encrypt(self,message):
        with open(os.path.join(basedir,self.public_file_name)) as f:
            key = f.read()
            rsakey = RSA.importKey(str(key))
            cipher = Cipher_pkcs1_v1_5.new(rsakey)
            cipher_text = base64.b64encode(cipher.encrypt(bytes(message.encode("utf8"))))
            return cipher_text
    
    def decrypt(self,cipher_text):
        with open(os.path.join(basedir,self.private_file_name)) as f:
            key = f.read()
            rsakey = RSA.importKey(key)
            cipher = Cipher_pkcs1_v1_5.new(rsakey)
            text = cipher.decrypt(base64.b64decode(cipher_text), self.random_generator)
            return str(text,"utf-8")


if __name__ == "__main__":
    pcras = RsaEncryption(message='123456',private_file_name='private1.pem',public_file_name='public2.pem')
    cipher_text=pcras.encrypt()
    print(cipher_text)
    print(pcras.decrypt(cipher_text))
